"""
视频水印去除工具 - 完整版
支持交互式鼠标选择和多种去除方法
"""

import cv2
import numpy as np
from typing import Tuple, Optional
import os
import sys

# 确保控制台中文显示正常
if os.name == 'nt':  # Windows系统
    import locale
    try:
        locale.setlocale(locale.LC_ALL, 'Chinese_China.utf8')
    except:
        pass


# ==================== 窗口管理 ====================

WINDOW_TITLES = {
    "watermark_selector": "Watermark Selector",
    "video_frame": "Video Frame Preview", 
    "comparison": "Before vs After"
}

def get_window_title(key: str) -> str:
    """获取窗口标题"""
    return WINDOW_TITLES.get(key, "OpenCV Window")


# ==================== 核心水印去除功能 ====================

def remove_watermark_by_inpaint(input_path: str, output_path: str, 
                               watermark_region: Tuple[int, int, int, int]) -> bool:
    """通过图像修复的方式去除水印"""
    try:
        cap = cv2.VideoCapture(input_path)
        if not cap.isOpened():
            print(f"无法打开视频文件: {input_path}")
            return False
            
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        x, y, w, h = watermark_region
        mask = np.zeros((height, width), dtype=np.uint8)
        mask[y:y+h, x:x+w] = 255
        
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        
        frame_count = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
                
            inpainted_frame = cv2.inpaint(frame, mask, 3, cv2.INPAINT_TELEA)
            out.write(inpainted_frame)
            
            frame_count += 1
            if frame_count % 30 == 0:
                print(f"处理进度: {frame_count}/{total_frames}")
                
        cap.release()
        out.release()
        print(f"视频处理完成，保存至: {output_path}")
        return True
        
    except Exception as e:
        print(f"处理视频时出错: {str(e)}")
        return False


def remove_watermark_by_blur(input_path: str, output_path: str,
                           watermark_region: Tuple[int, int, int, int],
                           blur_strength: int = 15) -> bool:
    """通过模糊处理的方式去除水印"""
    try:
        cap = cv2.VideoCapture(input_path)
        if not cap.isOpened():
            return False
            
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
        
        x, y, w, h = watermark_region
        frame_count = 0
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
                
            watermark_area = frame[y:y+h, x:x+w]
            blurred_area = cv2.GaussianBlur(watermark_area, (blur_strength, blur_strength), 0)
            frame[y:y+h, x:x+w] = blurred_area
            out.write(frame)
            
            frame_count += 1
            if frame_count % 30 == 0:
                print(f"处理进度: {frame_count}/{total_frames}")
                
        cap.release()
        out.release()
        print(f"视频处理完成，保存至: {output_path}")
        return True
        
    except Exception as e:
        print(f"处理视频时出错: {str(e)}")
        return False


def quick_remove_watermark(input_video: str, x: int, y: int, w: int, h: int, method: str = "inpaint"):
    """快速去除水印的简化函数"""
    output_video = f"无水印_{input_video}"
    
    if method == "inpaint":
        success = remove_watermark_by_inpaint(input_video, output_video, (x, y, w, h))
    elif method == "blur":
        success = remove_watermark_by_blur(input_video, output_video, (x, y, w, h))
    else:
        print(f"不支持的方法: {method}")
        return False
    
    if success:
        print(f"✅ 成功生成无水印视频: {output_video}")
    else:
        print("❌ 处理失败")
    
    return success


# ==================== 交互式选择器 ====================

class WatermarkSelector:
    def __init__(self):
        self.drawing = False
        self.start_point = None
        self.end_point = None
        self.original_frame = None
        self.display_frame = None
        self.scale = 1.0
        self.window_name = get_window_title("watermark_selector")
        
    def mouse_callback(self, event, x, y, flags, param):
        """鼠标回调函数"""
        if event == cv2.EVENT_LBUTTONDOWN:
            self.drawing = True
            self.start_point = (x, y)
            self.end_point = (x, y)
        elif event == cv2.EVENT_MOUSEMOVE:
            if self.drawing:
                self.end_point = (x, y)
                self.update_display()
        elif event == cv2.EVENT_LBUTTONUP:
            self.drawing = False
            self.end_point = (x, y)
            self.update_display()
            
    def update_display(self):
        """更新显示图像"""
        temp_frame = self.display_frame.copy()
        
        if self.start_point and self.end_point:
            cv2.rectangle(temp_frame, self.start_point, self.end_point, (0, 255, 0), 2)
            
            if self.start_point != self.end_point:
                x1 = int(min(self.start_point[0], self.end_point[0]) / self.scale)
                y1 = int(min(self.start_point[1], self.end_point[1]) / self.scale)
                x2 = int(max(self.start_point[0], self.end_point[0]) / self.scale)
                y2 = int(max(self.start_point[1], self.end_point[1]) / self.scale)
                
                width = x2 - x1
                height = y2 - y1
                
                info_text = f"({x1},{y1}) {width}x{height}"
                text_x = self.start_point[0]
                text_y = self.start_point[1] - 10 if self.start_point[1] > 30 else self.start_point[1] + 25
                
                text_size = cv2.getTextSize(info_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
                cv2.rectangle(temp_frame, (text_x-2, text_y-20), (text_x + text_size[0] + 4, text_y + 5), (0, 0, 0), -1)
                cv2.putText(temp_frame, info_text, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
        
        cv2.imshow(self.window_name, temp_frame)
        
    def select_region(self, frame: np.ndarray) -> Optional[Tuple[int, int, int, int]]:
        """交互式选择水印区域"""
        self.original_frame = frame.copy()
        height, width = frame.shape[:2]
        
        # 智能缩放
        max_width = 1200
        max_height = 800
        
        scale_w = max_width / width if width > max_width else 1
        scale_h = max_height / height if height > max_height else 1
        self.scale = min(scale_w, scale_h)
        
        if self.scale < 1:
            new_width = int(width * self.scale)
            new_height = int(height * self.scale)
            self.display_frame = cv2.resize(frame, (new_width, new_height))
            print(f"显示尺寸: {new_width} x {new_height} (缩放比例: {self.scale:.2f})")
        else:
            self.display_frame = frame.copy()
            print("按原尺寸显示")
        
        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.resizeWindow(self.window_name, self.display_frame.shape[1], self.display_frame.shape[0])
        cv2.setMouseCallback(self.window_name, self.mouse_callback)
        
        print(f"\n原始视频尺寸: {width} x {height}")
        print("操作说明: 用鼠标拖动选择水印区域，Enter确认，ESC取消，R重新选择")
        
        cv2.imshow(self.window_name, self.display_frame)
        
        while True:
            key = cv2.waitKey(1) & 0xFF
            
            if key == 13:  # Enter键
                if self.start_point and self.end_point and self.start_point != self.end_point:
                    x1 = int(min(self.start_point[0], self.end_point[0]) / self.scale)
                    y1 = int(min(self.start_point[1], self.end_point[1]) / self.scale)
                    x2 = int(max(self.start_point[0], self.end_point[0]) / self.scale)
                    y2 = int(max(self.start_point[1], self.end_point[1]) / self.scale)
                    
                    x1 = max(0, min(x1, width-1))
                    y1 = max(0, min(y1, height-1))
                    x2 = max(0, min(x2, width-1))
                    y2 = max(0, min(y2, height-1))
                    
                    region_width = x2 - x1
                    region_height = y2 - y1
                    
                    if region_width > 10 and region_height > 10:
                        cv2.destroyAllWindows()
                        return (x1, y1, region_width, region_height)
                    else:
                        print("选择区域太小，请重新选择")
                else:
                    print("请先选择一个区域")
                    
            elif key == 27:  # ESC键
                cv2.destroyAllWindows()
                return None
                
            elif key == ord('r') or key == ord('R'):  # R键重新选择
                self.start_point = None
                self.end_point = None
                self.drawing = False
                cv2.imshow(self.window_name, self.display_frame)
                print("重新选择区域...")


# ==================== 主要界面函数 ====================

def interactive_remove_watermark():
    """交互式水印去除主函数"""
    video_files = [f for f in os.listdir('.') if f.endswith(('.mp4', '.avi', '.mov', '.mkv'))
                   and not f.startswith('无水印_')]
    
    if not video_files:
        print("当前目录中没有找到视频文件")
        return
    
    print("=== 交互式视频水印去除工具 ===")
    print("\n找到的视频文件:")
    for i, video in enumerate(video_files):
        print(f"{i+1}. {video}")
    
    try:
        choice = int(input(f"\n请选择要处理的视频 (1-{len(video_files)}): ")) - 1
        if choice < 0 or choice >= len(video_files):
            print("选择无效")
            return
        input_video = video_files[choice]
    except ValueError:
        print("输入无效")
        return
    
    print(f"\n选择的视频: {input_video}")
    
    # 读取视频第一帧
    cap = cv2.VideoCapture(input_video)
    ret, first_frame = cap.read()
    cap.release()
    
    if not ret:
        print("无法读取视频文件")
        return
    
    # 交互式选择水印区域
    selector = WatermarkSelector()
    watermark_region = selector.select_region(first_frame)
    
    if watermark_region is None:
        print("取消选择")
        return
    
    x, y, w, h = watermark_region
    print(f"\n选择的水印区域: x={x}, y={y}, width={w}, height={h}")
    
    # 选择处理方法
    print("\n选择处理方法:")
    print("1. 图像修复 (推荐)")
    print("2. 模糊处理")
    
    try:
        method_choice = int(input("请选择方法 (1-2): "))
    except ValueError:
        method_choice = 1
    
    output_video = f"无水印_{input_video}"
    print(f"\n开始处理视频...")
    
    if method_choice == 2:
        success = remove_watermark_by_blur(input_video, output_video, watermark_region, 21)
    else:
        success = remove_watermark_by_inpaint(input_video, output_video, watermark_region)
    
    if success:
        print(f"\n✅ 处理完成！输出文件: {output_video}")
    else:
        print("\n❌ 处理失败")


def main():
    """主函数"""
    print("=== 视频水印去除工具 ===")
    print("1. 交互式鼠标选择 (推荐)")
    print("2. 快速处理 (已知坐标)")
    
    try:
        choice = int(input("\n请选择模式 (1-2): "))
    except ValueError:
        choice = 1
    
    if choice == 1:
        interactive_remove_watermark()
    elif choice == 2:
        print("\n快速处理示例:")
        print("quick_remove_watermark('video.mp4', x=100, y=100, w=200, h=100, method='inpaint')")
        print("\n请在Python中直接调用 quick_remove_watermark 函数")
    else:
        print("选择无效")


if __name__ == "__main__":
    main()